package com.blue.controller.ai;

import com.blue.common.entity.ResponseResult;
import com.blue.dto.ai.ChatRequestDTO;
import com.blue.service.system.IFileService;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.ai.model.Media;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiImageModel;
import org.springframework.ai.openai.OpenAiImageOptions;
import org.springframework.ai.openai.api.OpenAiImageApi;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.http.MediaType;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.net.URLConnection;
import java.nio.charset.StandardCharsets;
import java.util.Base64;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

@RestController
@RequestMapping("/ai/openAi")
public class OpenAiController {
    @Autowired
    private OpenAiImageModel openaiImageModel;
    @Autowired
    private IFileService fileService;
    @Autowired
    private OpenAiChatModel openAiChatModel;
    @Value("${file.uploadPath}")
    private String uploadPath;
    private final ChatMemory chatMemory = new InMemoryChatMemory();
    private final ChatClient chatClient;

    public OpenAiController(ChatClient.Builder chatClientBuilder) {
        this.chatClient = chatClientBuilder.build();
    }

    /**
     * 对话
     *
     * @param requestDTO
     * @return
     */
    @PostMapping("/chat")
    public String chat(@RequestBody ChatRequestDTO requestDTO) {
        return this.chatClient.prompt()
                .user(requestDTO.getMessage())
                .call()
                .content();
    }

    /**
     * 流式对话
     *
     * @param requestDTO
     * @return
     */
    @PostMapping(value = "/streamChat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> streamChat(@RequestBody ChatRequestDTO requestDTO) {
        return this.chatClient.prompt()
                .user(requestDTO.getMessage())
                .stream()
                .content();
    }

    /**
     * 上下文流式对话
     *
     * @param requestDTO
     * @return
     */
    @PostMapping(value = "/streamContextChat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> streamContextChat(@RequestBody ChatRequestDTO requestDTO) {
        ChatClient chatClient = ChatClient.builder(openAiChatModel)
                .defaultAdvisors(new PromptChatMemoryAdvisor(chatMemory))
                .build();
        return chatClient.prompt()
                .user(requestDTO.getMessage())
                .advisors(item -> item
                        .param(CHAT_MEMORY_CONVERSATION_ID_KEY, requestDTO.getContextId())
                        //历史会话最多10条发给AI
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10)
                )
                .stream()
                .content();
    }

    /**
     * 文生图
     *
     * @param requestDTO
     * @return
     */
    @PostMapping("/textToImg")
    public ResponseResult<?> chat1(@RequestBody ChatRequestDTO requestDTO) throws IOException {
        ImageResponse response = openaiImageModel.call(
                new ImagePrompt(requestDTO.getMessage(),
                        OpenAiImageOptions.builder()
                                //让图片更清晰
                                .withQuality("hd")
                                //使用文生图模型
                                .withModel(OpenAiImageApi.DEFAULT_IMAGE_MODEL)
                                //生成几张
                                .withN(1)
                                //图片尺寸
                                .withHeight(1024)
                                .withWidth(1024).build())

        );
        String url = response.getResult().getOutput().getUrl();
        InputStream inputStream = urlToStream(url);
        String code = fileService.uploadFile(inputStream);
        return ResponseResult.okResult(code);
    }

    private InputStream urlToStream(String imageUrl) throws IOException {
        // 将在线图片地址转换为URL对象
        URL url = new URL(imageUrl);
        // 打开URL连接
        URLConnection connection = url.openConnection();
        // 转换为HttpURLConnection对象
        HttpURLConnection httpURLConnection = (HttpURLConnection) connection;
        // 获取输入流
        return httpURLConnection.getInputStream();
    }

    /**
     * 多模态
     *
     * @param requestDTO
     * @return
     */
    @PostMapping(value = "/multiChat", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> multiChat(@RequestBody ChatRequestDTO requestDTO) {
        String fileName = new String(Base64.getDecoder().decode(requestDTO.getCode()), StandardCharsets.UTF_8);
        File file = new File(uploadPath + fileName);
        Resource resource = new FileSystemResource(file);
        return ChatClient.create(openAiChatModel)
                .prompt()
                .user(u -> u.text(requestDTO.getMessage())
                        .media(new Media(MimeTypeUtils.IMAGE_PNG, resource)))
                .stream()
                .content();
    }
}

